#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <arpa/inet.h>
#include <fcntl.h>
#include <unistd.h>
#include <libgen.h>
#include <sys/stat.h>
#include <sys/select.h>

#define PORT 3456//服务器监听端口
#define SERVER_IP "127.0.0.1"
#define BLOCK_SIZE 512//TFTP数据块大小
#define BASE_DIR "/home/linux/tftp_files"//文件存储根目录

typedef enum{
    //读请求
    OP_RRQ = 1,
    //写请求
    OP_WRQ = 2,
    //数据包
    OP_DATA = 3,
    //确认包
    OP_ACK = 4,
    //错误包
    OP_ERROR = 5
}tftp_opcode;

//发送请求函数
void send_request(int sockfd, struct sockaddr_in *addr, char *filename, int is_write){
    //初始化缓冲区
    uint8_t buf[516] = {0};
    buf[1] = is_write ? OP_WRQ:OP_RRQ;//设置操作码 is_write :1表示上传WRQ；0表示下载RRQ

    //写入文件名(从第三个字节开始，第1，2字节为操作码)
    strcpy((char *)(buf + 2), filename);

    //写入模式（octet表示二进制传输，从文件名后+1字节开始）
    strcpy((char*)(buf + 2 +strlen(filename)+ 1), "octet");

    //计算数据包长度：操作码（2） + 文件名（len） + 0 + 模式（len） + 0
    int len = 2 + strlen(filename) + 1 + strlen("octet") + 1;
    sendto(sockfd, buf, len, 0, (struct sockaddr*)addr, sizeof(*addr));
}

//文件下载函数
void download_file(int sockfd, struct sockaddr_in *addr, char *remote_filename, char *local_path){
    struct stat st;//用于存储本地路径状态

    //检查本地路径是否存在
    if(stat(local_path, &st) == 0 && S_ISDIR(st.st_mode)){
        //若是目录，使用远程文件名作为本地文件名
        char *filename = basename(remote_filename);
        char filepath[256];
        snprintf(filepath, sizeof(filepath), "%s/%s", local_path, filename);
        local_path = filepath;//更新为完整的本地文件路径
    }
    printf("下载到本地路径: %s\n", local_path);
    //创建或者打开本地文件
    int fd = open(local_path, O_WRONLY | O_CREAT | O_TRUNC, 0644);
    if(fd < 0){
        perror("无法创建本地文件");
        return;
    }
    uint8_t buf[BLOCK_SIZE + 4];//数据包缓冲区（512字节数据+4字节头部）
    uint16_t expected_block = 1;//期望的数据块编号从1开始

    while(1){
        //服务器地址长度
        socklen_t len = sizeof(*addr);
        //接收数据包(包含服务器地址)
        ssize_t nbytes = recvfrom(sockfd, buf, sizeof(buf), 0, (struct sockaddr*)addr, &len);

        if(nbytes < 4){//数据包过小
            printf("收到无效数据包，忽略\n");
            break;
        }
        //处理数据块
        if(buf[1] == OP_DATA && ntohs(*(uint16_t*)&buf[2]) == expected_block){
            write(fd, &buf[4], nbytes - 4);//写入文件(跳过四个字节包头)
            //构建ACK包（直接使用接收到的块号字节）
            uint8_t ack[4] = {0, OP_ACK, buf[2], buf[3]};
            sendto(sockfd, ack, sizeof(ack), 0, (struct sockaddr*)addr, sizeof(*addr));//发送ACK包
            //如果数据库不足512字节，视为最后一块
            if(nbytes - 4 < BLOCK_SIZE){
                break;
            }
            expected_block++;//期望块号递增
        }else if(buf[1] == OP_ERROR){
            //处理错误包
            printf("错误： %s\n", &buf[4]);
            break;
        }
    }
    close(fd);//关闭文件描述符
}

void upload_file(int sockfd, struct sockaddr_in *addr, char *local_path, char *remote_filename){
    //以只读模式打开本地文件
    int fd = open(local_path, O_RDONLY);
    if(fd < 0){
        perror("open failed:");
        return;
    }

    //发送数据块缓冲区
    uint8_t buf[BLOCK_SIZE + 4];
    uint16_t block = 1;
    ssize_t bytes_read;

    //等待服务器响应WRQ的初始ACK，块号0
    recvfrom(sockfd, buf, 4, 00, NULL, NULL);

    while((bytes_read = read(fd, buf+4, BLOCK_SIZE) > 0)){
        //构建DATA包头
        buf[0] = 0;
        buf[1] = OP_DATA;
        buf[2] = block >> 8;
        buf[3] = block && 0xFF;
        //发送数据块
        sendto(sockfd, buf,bytes_read + 4, 0, (struct sockaddr*)addr, sizeof(*addr));

        //等待服务器响应
        recvfrom(sockfd, buf, 4, 0, NULL, NULL);
        block++;
    }
    close(fd);
}
/*  下载：
    1，客户端发送RRQ操作码1+文件名+octet到服务器中
    2，服务器解析RRQ，检测文件存在性
    存在：分块发送DATA包，每发送一块等待ACK回应
    不存在：返回ERROR包
    3，客户端接收DATA包，验证块号后写入文件，返回对应ACK
    4，服务器发送空DATA包表示结束，客户端发送最后一个ACK

    上传：
    1，客户端发送“WRQ",
    2，服务器解析WRQ创建文件并返回ACK0
    3，客户端”接收ACK0“，分块发送DATA包，块号从1开始每发送一块等待ACK
    4，服务器接收DATA包，验证块号写入文件，然后发送ACK
    5，客户端发送完所有数据，服务器接收后结束
*/
int main(int argc, char *argv[]){

    if(argc != 6){
        printf("用法:%s <get|put> <服务器IP> <服务器端口> <远程文件名> <服务端ip地址> <端口号>\n", argv[0]);
        printf("示例: \n 下载：%s get /local/path remote.txt 192.169.1.100 69\n", argv[0]);
        return 1;
    }
    //创建UDPsocket
    int sockfd = socket(AF_INET, SOCK_DGRAM, 0);
    struct sockaddr_in server_addr;
    server_addr.sin_family = AF_INET;
    server_addr.sin_port = htons(atoi(argv[5]));
    server_addr.sin_addr.s_addr = inet_addr(argv[4]);
    //打印操作信息
    printf("操作: %s\n", argv[1]);
    printf("本地路径：%s\n", argv[2]);
    printf("远程文件名：%s\n", argv[3]);

    //根据操作类型执行上传或者下载
    if(strcmp(argv[1], "get") == 0){
        //下载文件
        printf("开始下载文件...\n");
        //调用下载函数（实现略）
        send_request(sockfd, &server_addr, argv[3], 0);
        //接收数据
        download_file(sockfd, &server_addr, argv[3], argv[2]);
    }else if(strcmp(argv[1], "put") == 0){
        //上传文件
        printf("开始上传文件...\n");
        //调用上传函数（实现略）
        send_request(sockfd, &server_addr, argv[3], 1);
        upload_file(sockfd, &server_addr, argv[2], argv[3]);
    }else{
        printf("未知操作: %s\n", argv[1]);
        close(sockfd);
        return 1;
    }
    close(sockfd);
    return 0;
}